package com.tuniu.agents.agent;

import jakarta.annotation.PostConstruct;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.util.Assert;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * AgentManager implementation based on Spring
 */
public class SpringAgentManager implements ApplicationContextAware, AgentManager {

    private final ConcurrentHashMap<String, Agent> contextMap = new ConcurrentHashMap<>();

    private ApplicationContext applicationContext;

    @PostConstruct
    public void init() {
        // Initialize the contextMap with beans of type Agent from the application context
        for (Map.Entry<String, Agent> entry : applicationContext.getBeansOfType(Agent.class).entrySet()) {
            Agent agent = entry.getValue();
            contextMap.put(agent.getAgentId(), agent);
        }
    }

    @Override
    public Agent getAgent(String agentId) {
        Agent agent = contextMap.get(agentId);
        if (agent == null) {
            // If the agent is not found in the contextMap, retrieve it from the application context
            agent = applicationContext.getBean(agentId, Agent.class);
            if (agent != null) {
                contextMap.put(agentId, agent);
            }
        }
        // Ensure that the agent exists, throw an exception if it does not
        Assert.notNull(agent, "Agent not exists: " + agentId);
        return agent;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }
}
